#! /bin/bash

RED="\033[31m"
GREEN="\033[32m"
YELLOW="\033[33m"
BLUE="\033[34m"
END="\033[0m"

# root权限
root_privilege() {
    if [[ $EUID -ne 0 ]]; then
        echo "Error: This script must be run as root!" 1>&2
        exit 1
    fi
}
# 系统信息
system_info() {
    clear
    ## 定义系统名称(Define system name)
    SYSTEM_NAME=$(cat /etc/os-release  | grep -E "^NAME=" | awk -F '=' '{print$2}' | sed "s/[\'\"]//g")
    ## 定义系统版本(Define system version)
    SYSTEM_VERSION_NUMBER=$(cat /etc/os-release  | grep -E "VERSION_ID=" | awk -F '=' '{print$2}' | sed "s/[\'\"]//g")
    ARCH=$(uname -m)
}

# 允许ip
allow_ipset() {
    check_ipset
    # 添加ipset规则
    echo -e "${BLUE}查看国家代码： https://yumingsuoxie.bmcx.com 或者 https://www.ipdeny.com/ipblocks${END}"
    echo -e "${BLUE}请输入允许访问的国家代码，如中国(cn)、美国(us)、日本(jp)${END}"
    read -p "$(echo -e "\n${YELLOW}└─ Please enter the country code： ${END}")" COUNTRY_CODE
    COUNTRY_CODE=$(echo "$COUNTRY_CODE" | tr '[:upper:]' '[:lower:]')
    echo -ne "\n${BLUE} 正在下载IP数据源，请耐心等待...${END}\r"
    wget -P /tmp http://www.ipdeny.com/ipblocks/data/countries/$COUNTRY_CODE.zone 2> /dev/null
    # 检查下载是否成功
    if [ -f "/tmp/"$COUNTRY_CODE".zone" ]; then
        echo -ne "${GREEN}\033[K IP源下载完成！${END}\r"
    else
        echo -ne "${BLUE}\033[K 下载失败，请检查你的输入！${END}\r"
        echo -e "${BLUE} 查看下载地址：https://www.ipdeny.com/ipblocks${END}\n"
        exit 1
    fi
    # 创建本地可访问规则
    local_ips=$(iptables -L INPUT --line-numbers | grep local_ips | wc -l)
    if [ $local_ips -eq 0 ]; then
        ipset create local_ips hash:net
        ipset add local_ips 10.0.0.0/8
        ipset add local_ips 172.16.0.0/12
        ipset add local_ips 192.168.0.0/16
        iptables -I INPUT -m set --match-set local_ips src -j ACCEPT
    fi
    # 创建其他规则
    ipset create $COUNTRY_CODE hash:net
    for i in $(cat /tmp/$COUNTRY_CODE.zone ); do ipset -A $COUNTRY_CODE $i; done
    rm -rf /tmp/$COUNTRY_CODE.zone*
    echo -ne "${GREEN}\033[K 正在添加IP的允许规则！${END}\r"
    # iptables允许规则，在头部插入
    iptables -I INPUT -m set --match-set "$COUNTRY_CODE" src -j ACCEPT
    # iptables禁止其他IP访问，在尾部追加
    drop_all=$(iptables -L INPUT --line-numbers | grep -P '^\d+\s+DROP\s+all\s+--\s+anywhere\s+anywhere\s*$' | wc -l)
    if [ $drop_all -eq 0 ]; then
        iptables -A INPUT -j DROP
    fi
    echo -e "${GREEN}\033[K 所指定国家($COUNTRY_CODE)的IP已授权访问！${END}\n"
}
# 封禁ip
block_ipset() {
    check_ipset
    # 添加ipset规则
    echo -e "${BLUE}查看国家代码： https://yumingsuoxie.bmcx.com 或者 https://www.ipdeny.com/ipblocks${END}"
    echo -e "${BLUE}请输禁止访问的国家代码，如如中国(cn)、美国(us)、日本(jp)${END}"
    read -p "$(echo -e "\n${YELLOW}└─ Please enter the country code： ${END}")" COUNTRY_CODE
    COUNTRY_CODE=$(echo "$COUNTRY_CODE" | tr '[:upper:]' '[:lower:]')
    echo -ne "\n${BLUE} 正在下载IP数据源，请耐心等待...${END}\r"
    wget -P /tmp http://www.ipdeny.com/ipblocks/data/countries/$COUNTRY_CODE.zone 2> /dev/null
    # 检查下载是否成功
    if [ -f "/tmp/"$COUNTRY_CODE".zone" ]; then
        echo -ne "${GREEN}\033[K IP源下载完成！${END}\r"
    else
        echo -ne "${BLUE}\033[K 下载失败，请检查你的输入！${END}\r"
        echo -e "${BLUE} 查看下载地址：https://www.ipdeny.com/ipblocks${END}\n"
        exit 1
    fi
    # 创建规则
    ipset create $COUNTRY_CODE hash:net
    for i in $(cat /tmp/$COUNTRY_CODE.zone ); do ipset -A $COUNTRY_CODE $i; done
    rm -rf /tmp/$COUNTRY_CODE.zone*
    echo -ne "${GREEN}\033[K 正在添加IP的禁止规则！${END}\r"
    # 开始封禁
    iptables -I INPUT -m set --match-set "$COUNTRY_CODE" src -j DROP
    echo -e "${GREEN}\033[K 所指定国家($COUNTRY_CODE)的IP已禁止访问！${END}\n"
}

# 删除iptables规则
unblock_ipset() {
    echo ""
    iptables -L INPUT --line-numbers | awk '/match-set/ {print $1, $8, $2}' | sed -E 's/(.*)/\x1b[34m\1\x1b[0m/'
    iptables -L INPUT --line-numbers | grep -P '^\d+\s+DROP\s+all\s+--\s+anywhere\s+anywhere\s*$' | sed -E 's/(.*)/\x1b[34m\1\x1b[0m/'
    echo -e "\n${BLUE}请输入规则列表中的需要删除的编号${END}"
    read -p "$(echo -e "\n${YELLOW}└─ Please enter the number： ${END}")" NUM_CODE

    iptables_list=$(iptables -L INPUT --line-numbers | grep 'match-set' | wc -l)
    
    if [[ $NUM_CODE =~ ^[0-9]+$ ]] && [[ $NUM_CODE -le $((iptables_list + 1)) ]]; then
       iptables_name=$(iptables -L INPUT --line-numbers | grep "^$NUM_CODE" | awk '/match-set/ {print $8}')
        if [ "$iptables_name" = "local_ips" ]; then
            echo -e "\n${RED}局域网内部规则不允许被删除！${END}\n"
            exit 1
        elif [ -z "$iptables_name" ]; then
            iptables -D INPUT $NUM_CODE
            echo ""
        else
            iptables -D INPUT $NUM_CODE
            ipset destroy $iptables_name
            echo ""
        fi
    else 
        echo -e "\n${RED} 输入有误，请重新输入！${END}\n"
        unblock_ipset
    fi
}

block_city() {
    check_ipset
    # 创建本地局域网
    local_ips=$(iptables -L INPUT --line-numbers | grep local_ips | wc -l)
    if [ $local_ips -eq 0 ]; then
        echo -ne "${YELLOW}\033[K 局域网IP，正在添加中...${END}\r"
        ipset create local_ips hash:net
        ipset add local_ips 10.0.0.0/8
        ipset add local_ips 172.16.0.0/12
        ipset add local_ips 192.168.0.0/16
        iptables -I INPUT -m set --match-set local_ips src -j ACCEPT
        echo -ne "${GREEN}\033[K 局域网IP，完成！${END}\r"
    fi

    sed -i 's/\r$//' IPs_*
    # 重庆
    cq=$(ipset list chongqing 2>/dev/null | wc -l)
    if [ $cq -eq 0 ]; then
        echo -ne "${YELLOW}\033[K 重庆IP，正在添加中...${END}\r"
        ipset create chongqing hash:net
        for ip_ranges in $(cat IPs_Chongqing.txt); do ipset add chongqing $ip_ranges ;done
    fi
    iptables_cq=$(iptables -L --line-numbers | grep match-set | grep chongqing | wc -l)
    if [ $iptables_cq -eq 0 ]; then
        iptables -I INPUT -m set --match-set chongqing src -j ACCEPT
        echo -e "${GREEN}\033[K 重庆IP，完成！${END}\r"
    fi
    # 上海
    echo ""
    sh=$(ipset list shanghai 2>/dev/null | wc -l)
    if [ $sh -eq 0 ]; then
        echo -ne "${YELLOW}\033[K 上海IP，正在添加中...${END}\r"
        ipset create shanghai hash:net 
        for ip_ranges in $(cat IPs_Shanghai.txt); do ipset add shanghai $ip_ranges ;done
    fi
    iptables_sh=$(iptables -L --line-numbers | grep match-set | grep shanghai | wc -l)
    if [ $iptables_sh -eq 0 ]; then
        iptables -I INPUT -m set --match-set shanghai src -j ACCEPT
        echo -e "${GREEN}\033[K 上海IP，完成！${END}\r"
    fi
    # iptables禁止其他IP访问，在尾部追加
    drop_all=$(iptables -L INPUT --line-numbers | grep -P '^\d+\s+DROP\s+all\s+--\s+anywhere\s+anywhere\s*$' | wc -l)
    if [ $drop_all -eq 0 ]; then
        iptables -A INPUT -j DROP
    fi
    echo ""
}

# 查看规则列表
block_list() {
    echo ""
    iptables -L --line-numbers | grep match-set
    iptables -L INPUT --line-numbers | grep -P '^\d+\s+DROP\s+all\s+--\s+anywhere\s+anywhere\s*$'
    echo ""
}

#检查ipset是否安装
check_ipset() {
    if [ -f /sbin/ipset ]; then
        echo -e ""
    elif [ $(echo SYSTEM_VERSION_NUMBER | grep Linux | wc -l) -ge 1 ]; then
        yum -y install ipset
    else
        apt-get -y install ipset
    fi
    firewalld_1=$(systemctl status firewalld.service 2>/dev/null | grep running | wc -l)
    if [ $firewalld_1 -eq 1 ]; then
        systemctl stop firewalld && systemctl disable firewalld &>/dev/null
        iptables -F && iptables -X && iptables -t nat -F && iptables -t nat -X
    fi
}

#开始菜单
main() {
root_privilege
system_info
echo -e ""
echo -e "+---------------------------------------------------+"
echo -e "|        \033[1m欢迎使用 IP 访问控制系统${END}                   |"
echo -e '+---------------------------------------------------+'
echo -e "         运行环境: ${BLUE}${SYSTEM_NAME} ${SYSTEM_VERSION_NUMBER}${END} "
echo -e "         系统时间: ${BLUE}$(date "+%Y-%m-%d %H:%M:%S")${END}            "
echo -e "+---------------------------------------------------+"
echo -e ""
echo -e "\033[36m         ❖  允许某国家(IP) \t  \t 1 )${END}"
echo -e "\033[36m         ❖  禁止某国家(IP) \t  \t 2 )${END}"
echo -e "\033[36m         ❖  删除允许/禁制规则 \t  \t 3 )${END}"
echo -e "\033[36m         ❖  查看规则列表 \t \t 4 )${END}"
echo -e "\033[36m         ❖  定制化城市(IP) \t \t 5 )${END}"

CHOICE_NUM=$(echo -e "\n${YELLOW}└─ Please enter the number? :  ${END}")
read -p "${CHOICE_NUM}" INPUT

case "$INPUT" in
    1)
    allow_ipset
    ;;
    2)
    block_ipset
    ;;
    3)
    unblock_ipset
    ;;
    4)
    block_list
    ;;
    5)
    block_city
    ;;
    *)
    echo -e "${YELLOW}└─ Please enter the correct number： ${END}"
    main
    ;;
    esac
}
main
